#!/usr/bin/env python
# coding=utf-8
# CopyRight(c) 2019 Inspur Electronic Information Industry Co.,Ltd

# Supported key sizes:
#   256-bit

# Supported modes of operation:
#   CBC - Cipher-Block Chaining

# Supported modes of padding:
# PKCS7 
# if length of plaintext <16 generate 16 bytes cipher_text
# if length of plaintext =16 padding 16bytes,and generate 32 bytes cipher_text

import sys
import getpass
import base64
from hashlib import pbkdf2_hmac

from Crypto.Cipher import AES
from Crypto.Hash import HMAC, SHA256
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad
from Crypto.Util.Padding import unpad

import check
import seccheck_log
import utils

AES_KEY_SIZE = 32
HMAC_KEY_SIZE = 16
IV_SIZE = 16
SALT_SIZE = 16
HMAC_SIZE = 32


def get_key_iv(password, salt, workload=10000):
    """
    generate an AES key, an hmac key and an AES
    initialization vector according password and salt
    Args:
        password: 8bytes init password
        salt: 16 bytes random salt value
        workload: Number of Iterations,default 10000
    Returns:
        AES key, an hmac key and an AES initialization vector
    Raises:
        None
    """
    stretched = pbkdf2_hmac('sha256', password, salt, workload, AES_KEY_SIZE +
                            IV_SIZE + HMAC_KEY_SIZE)
    aes_key, stretched = stretched[:AES_KEY_SIZE], stretched[AES_KEY_SIZE:]
    hmac_key, stretched = stretched[:HMAC_KEY_SIZE], stretched[HMAC_KEY_SIZE:]
    iv = stretched[:IV_SIZE]
    return aes_key, hmac_key, iv


def encrypt(pass_wd, plain_text, workload=1000):
    """
    Encrypts `plaintext` with `pass_wd` using AES-CBC-256, an hmac to verify
    integrity,and PBKDF2 to stretch the given pass_wd.
    Args:
        pass_wd: 8bytes init password
        plain_text: to be encrypted string
        workload: Number of Iterations,default 1000
    Returns:
        cipher_text,special include mac, salt and cipher_text
    Raises:
        ValueError:calc mac failed
    """
    if isinstance(pass_wd, str):
        pass_wd = pass_wd.encode('utf-8')
    if isinstance(plain_text, str):
        plain_text = plain_text.encode('utf-8')

    salt = get_random_bytes(SALT_SIZE)
    key, hmac_key, iv = get_key_iv(pass_wd, salt, workload)

    cipher = AES.new(key, AES.MODE_CBC, iv)
    cipher_text = cipher.encrypt(pad(plain_text, AES.block_size))
    h = HMAC.new(hmac_key, digestmod=SHA256)
    h.update(salt + cipher_text)
    hmac = h.digest()
    if len(hmac) != HMAC_SIZE:
        raise ValueError("calc mac failed")
    return hmac + salt + cipher_text


def decrypt(pass_wd, cipher_text, workload=1000):
    """
    Decrypts `cipher_text` with `pass_wd` using AES-CBC-256, an HMAC to verify
    integrity,and PBKDF2 to stretch the given pass_wd.
    Args:
        pass_wd: 8 bytes init password
        cipher_text: to be decrypted string
        workload: Number of Iterations,default 1000
    Returns:
        decrypted plaintext
    Raises:
        ValueError:check length or mac failed
    """

    if len(cipher_text) % 16 != 0:
        raise ValueError("Cipher_text must be made of full 16-byte blocks.")
    if len(cipher_text) < 64:
        raise ValueError("Cipher_text must be at least 64 bytes long.")
    if isinstance(pass_wd, str):
        pass_wd = pass_wd.encode('utf-8')

    hmac, cipher_text = cipher_text[:HMAC_SIZE], cipher_text[HMAC_SIZE:]
    salt, cipher_text = cipher_text[:SALT_SIZE], cipher_text[SALT_SIZE:]
    ase_key, hmac_key, iv = get_key_iv(pass_wd, salt, workload)
    h = HMAC.new(hmac_key, digestmod=SHA256)
    h.update(salt + cipher_text)
    try:
        h.verify(hmac)
    except ValueError:
        raise ValueError("Cipher_text corrupted or tampered or the key is wrong")
    encryptor = AES.new(ase_key, AES.MODE_CBC, iv)
    return unpad(encryptor.decrypt(cipher_text), AES.block_size)


def aes_decrypt(key, cipher_text, out=True):
    """decrypt cipher text, print and return plain text
    Args:
        key: key seed for generating secret key
        cipher_text: ciphertext to be deciphered
        out: True print console  False: not print
    Returns:
        plain_text: returns CBC dencrypted plaintext
    Raises:
        ValueError:
        base64.binascii.Error:
    """

    try:
        cipher_text_decode = base64.b64decode(cipher_text)
        text = decrypt(key, cipher_text_decode).decode('utf8')
        if out:
            print('decrypt cipher text succeeded:{}'.format(text))
        return text
    except (TypeError, base64.binascii.Error, ValueError) as e:
        raise e


def aes_encrypt(key, plain_text, out=True):
    """encrypt plain text, print and return cipher text
    Args:
        key: key seed for generating secret key
        plain_text: cleartext to be encrypted
        out: True print console  False: not print
    Returns:
        cipher_text: returns CBC encrypted ciphertext
    Raises:
        ValueError:
    """
    try:
        cipher_text = encrypt(key, plain_text)
        cipher_text_base64 = base64.b64encode(cipher_text)
        if out:
            print('encrypt plain text succeeded:{}'.format(cipher_text_base64.decode('utf8')))
        return cipher_text_base64
    except ValueError as e:
        print('encrypt occur some errors: {}'.format(e))
        raise e


def decrypt_plugin_config(plugin_info, key):
    """decrypt plugin config parameter, it may be encrypt
        Args:
            plugin_info: all plugin info
            key: key seed for generating secret key
        Returns:
            plugin_info: all plugin info,if plugin_config parameter is encrypted,returns will be decrypted
        Raises:
            BaseException:  decrypted fail
        """
    if not check.check_not_none_dict(plugin_info):
        return
    text = plugin_info.get('plugin_config', None)
    encry_flag = plugin_info.get('encryption', None)
    if encry_flag == 1:
        try:
            text = aes_decrypt(key, text, out=False)
            plugin_info['plugin_config'] = text
        except (base64.binascii.Error, ValueError) as err:
            seccheck_log.sec_check_logger.error('decrypt plugin config occur error:{}'.format(err))
            raise err
    seccheck_log.sec_check_logger.info('plugin_exe:{} plugin_config:{}'.format(plugin_info['plugin_exe'], text))


def aes_enc_decrypt(args, key):
    """
    encrypt or decrypt for AES_KEY
    Args:
        args: argument_list
        key: key seed for generating secret key
    Raises:
        ValueError: An error in encrypt ASE_KEY
    """

    # 处理字符串加密
    if args.get('encrypt', None):
        try:
            aes_encrypt(key, args.get('encrypt', None))
        except ValueError as err:
            print('encrypt AES_KEY occur some errors: {}'.format(err))
        finally:
            utils.program_exit(sys.argv[0], 0)

    # 处理字符串解密
    if args.get('decrypt', None):
        try:
            aes_decrypt(key, args.get('decrypt', None))
        except (TypeError, base64.binascii.Error, ValueError) as err:
            print('decrypt AES_KEY occur some errors: {}'.format(err))
        finally:
            utils.program_exit(sys.argv[0], 0)


def get_input_secret_key(python_version, secret_key):
    """
    get secret key from terminal
    Args:
        python_version: python version info
        secret_key: secret key
    Returns:
        return secret key or exit
    """

    while True:
        if not secret_key:
            secret_key = getpass.getpass("Enter a secret key(8-16 digits, alphabets and special characters):")
        if not check.validate_input_key(secret_key):
            confirm = python_version.input_info(
                "secret key consists of 8-16 digits, alphabets and special characters,"
                "re-enter the key or quit(Y/N)")
            if confirm == 'y' or confirm == "Y":
                secret_key = getpass.getpass("Enter a secret key:")
                continue
            else:
                print('not input key,the process auto exit')
                utils.program_exit(sys.argv[0], 0)
        else:
            break
    return secret_key
